from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain_core.language_models import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.retrievers import RetrieverOutputLike
from langchain_core.vectorstores import VectorStoreRetriever

from config.llm_client import create_chat_client
from know_retriever import FaqRetriever


def create_context_retriever(llm: BaseChatModel, retriever: VectorStoreRetriever) -> RetrieverOutputLike:
    contextualize_q_system_prompt = """
    给定聊天历史和最新的用户问题（可能引用聊天历史中的上下文），
    制定一个独立的问题，该问题可以在没有聊天历史的情况下被理解。
    不要回答问题，只是根据需要重新制定问题，否则按原样返回。
    """
    contextualize_q_prompt = ChatPromptTemplate.from_messages([
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])
    # 用于感知会话历史的上下文，再根据上下文来检索相关文档，提高精确性
    return create_history_aware_retriever(
        llm, retriever, contextualize_q_prompt
    )


if __name__ == '__main__':
    retriever = FaqRetriever()
    documents = retriever.retriever.invoke("航班延误处理流程")
    for i, document in enumerate(documents, start=1):
        print(f"Document{i}:{document.page_content}+{document.metadata.get("answer")}")
    llm = create_chat_client()
    # create_context_retriever(llm,retriever.retriever)
